#include "Lattice.h"

const vec<double> e[Q_DIM] = {{0, 0}, {1.,0}, {0,1.}, {-1.,0}, {0,-1.}, {1.,1.}, {-1.,1.}, {-1.,-1.}, {1.,-1.}};
const int opositeDirection[Q_DIM] = {0, 3, 4, 1, 2, 7, 8, 5, 6};
const double w[Q_DIM] = {4./9., 1./9., 1./9., 1./9., 1./9., 1./36., 1./36., 1./36., 1./36.};

const double M[Q_DIM][Q_DIM] = { 	{ 1, 1, 1, 1, 1, 1, 1, 1, 1},
									{-4,-1,-1,-1,-1, 2, 2, 2, 2},
									{ 4,-2,-2,-2,-2, 1, 1, 1, 1},
									{ 0, 1, 0,-1, 0, 1,-1,-1, 1},
									{ 0,-2, 0, 2, 0, 1,-1,-1, 1},
									{ 0, 0, 1, 0,-1, 1, 1,-1,-1},
									{ 0, 0,-2, 0, 2, 1, 1,-1,-1},
									{ 0, 1,-1, 1,-1, 0, 0, 0, 0},
									{ 0, 0, 0, 0, 0, 1,-1, 1,-1}}; //M[i][j]

const double Minv[Q_DIM][Q_DIM] = {	{ 4,-4, 4, 0, 0, 0, 0, 0, 0},
									{ 4,-1,-2, 6,-6, 0, 0, 9, 0},
									{ 4,-1,-2, 0, 0, 6,-6,-9, 0},
									{ 4,-1,-2,-6, 6, 0, 0, 9, 0},
									{ 4,-1,-2, 0, 0,-6, 6,-9, 0},
									{ 4, 2, 1, 6, 3, 6, 3, 0, 9},
									{ 4, 2, 1,-6,-3, 6, 3, 0,-9},
									{ 4, 2, 1,-6,-3,-6,-3, 0, 9},
									{ 4, 2, 1, 6, 3,-6,-3, 0,-9}}; //36M^{-1}[i][j] (falta dividir por 36)

Lattice::Lattice() {
}

void Lattice::init(int nX, int nY, double _rho, double se, double sepsilon, double sq, double snu) {
	dim = {nX, nY};

	f = vec3double(dim.x, vec2double(dim.y, vecdouble(Q_DIM)));
	faux = vec3double(dim.x, vec2double(dim.y, vecdouble(Q_DIM)));
	m = vec3double(dim.x, vec2double(dim.y, vecdouble(Q_DIM)));
	//rho = vec2double(dim.x, vecdouble(dim.y));
	//u = vector<vector<vec<double> > >(dim.x, vector<vec<double> >(dim.y));
	isObstacle = vec2bool(dim.x, vector<bool>(dim.y));
	
	resetGrid(_rho);

	s[iE] = se;
	s[iEPSILON] = sepsilon;
	s[iQX] = sq;
	s[iQY] = sq;
	s[iPXX] = snu;
	s[iPXY] = snu;
}

void Lattice::resetGrid(double _rho) {
	for (int i = 0; i < dim.x; i++)
		for (int j = 0; j < dim.y; j++) {
			//rho[i][j] = _rho;
			//u[i][j] = {0.0, 0.0};
			isObstacle[i][j] = false;
			for (int l = 0; l < Q_DIM; l++) {
				f[i][j][l] = 0.0;
				faux[i][j][l] = 0.0;
				m[i][j][l] = 0.0;
			}
		}
}

void Lattice::swapGrids() {
	for (int i = 0; i < dim.x; i++)
		for (int j = 0; j < dim.y; j++)
			for (int l = 0; l < Q_DIM; l++)
				f[i][j][l] = faux[i][j][l];
}

void Lattice::stream(int i, int j) {
	for (int l = 0; l < Q_DIM; l++) {
		int _opdir = opositeDirection[l];
		vec<int> fromCell = {i + (int)e[_opdir].x, j + (int)e[_opdir].y};

		if (fromCell.x < 0)
			fromCell.x = dim.x-1;
		else if (fromCell.x == dim.x)
			fromCell.x = 0;

		if (fromCell.y < 0)
			fromCell.y = dim.y-1;
		else if (fromCell.y == dim.y)
			fromCell.y = 0;

		if (isObstacle[fromCell.x][fromCell.y])
			faux[i][j][l] = f[i][j][_opdir];
		else
			faux[i][j][l] = f[fromCell.x][fromCell.y][l];
	}
}

void Lattice::calcEQ(int i, int j) {
	meq[iE] = 0.25*alpha2*m[i][j][iRHO] + (1./6.)*gamma2*(m[i][j][iJX]*m[i][j][iJX] + m[i][j][iJY]*m[i][j][iJY]);
	meq[iEPSILON] = 0.25*alpha3*m[i][j][iRHO] + (1./6.)*gamma4*(m[i][j][iJX]*m[i][j][iJX] + m[i][j][iJY]*m[i][j][iJY]);
	meq[iQX] = 0.5*c1*m[i][j][iJX];
	meq[iQY] = 0.5*c1*m[i][j][iJY];
	meq[iPXX] = 1.5*gamma1*(m[i][j][iJX]*m[i][j][iJX] - m[i][j][iJY]*m[i][j][iJY]);
	meq[iPXY] = 1.5*gamma3*m[i][j][iJX]*m[i][j][iJY];
}

void Lattice::collide(int i, int j) {
	for(int ii=0; ii<Q_DIM; ii++) { //calculo los momentos 
		m[i][j][ii] = 0.0;
		for(int jj=0; jj<Q_DIM; jj++)
			m[i][j][ii] += M[ii][jj]*faux[i][j][jj];
	}

	calcEQ(i, j);

	for (int l=1; l<Q_DIM; l++)
		if ((l != iJX) && (l != iJY))
			m[i][j][l] -= s[l]*(m[i][j][l] - meq[l]);

	for(int ii=0; ii<Q_DIM; ii++) { //vuelvo a las DFs
		faux[i][j][ii] = 0.0;
		for(int jj=0; jj<Q_DIM; jj++)
			faux[i][j][ii] += (Minv[ii][jj]/36.)*m[i][j][jj];
	}
}

void Lattice::update() {
	for (int i = 0; i < dim.x; i++)
		for (int j = 0; j < dim.y; j++)
			if (isObstacle[i][j] == false) {
				stream(i, j);
				collide(i, j);
			}
	swapGrids();
}

double Lattice::energy() {
	double ret = 0.0;
	for(int i=0; i < dim.x; i++)
		for(int j=0; j < dim.y; j++) {
			double _ux = 0.0;
			double _uy = 0.0;
			double _rho = 0.0;
			for(int l=0; l < Q_DIM; l++) {
				_rho += f[i][j][l];
				_ux += M[iJX][l]*f[i][j][l];
				_uy += M[iJY][l]*f[i][j][l];
			}
			ret += (_ux*_ux + _uy*_uy);
		}
	return ret/2.0;
}

double Lattice::RMSvelocity() {
	return sqrt(2.*energy());
}

void Lattice::setU(int i, int j, vec<double> _u, double _rho) {
	if ((i >= 0) && (i < dim.x) && (j >= 0) && (j < dim.y) && (isObstacle[i][j] == false)) {
		m[i][j][iRHO] = _rho;
		m[i][j][iJX] = _u.x;
		m[i][j][iJY] = _u.y;
		calcEQ(i, j);
		for (int l=1; l<Q_DIM; l++)
			if ((l != iJX) && (l != iJY))
				m[i][j][l] = meq[l];

		for(int ii=0; ii<Q_DIM; ii++) {
			f[i][j][ii] = 0.0;
			for(int jj=0; jj<Q_DIM; jj++)
				f[i][j][ii] += (Minv[ii][jj]/36.)*m[i][j][jj];
		}
	}
}

double Lattice::getPopulation() {
	double ret = 0.0;
	for (int i = 0; i < dim.x; i++)
		for (int j = 0; j < dim.y; j++)
			if (isObstacle[i][j] == false) 
				for (int l = 0; l < Q_DIM; l++)
					ret += f[i][j][l];
	return ret;
}
